import math
from enum import Enum
from multiprocessing import Lock

import matplotlib.pyplot as plt
import numpy as np

from LaserData import get_obstacle_array
import time

def dwa_control(x, config, goal, ob):
    """
    Dynamic Window Approach control
    """

    dw = calc_dynamic_window(x, config)

    u, trajectory = calc_control_and_trajectory(x, dw, config, goal, ob)

    return u, trajectory


class RobotType:
    circle = 0
    rectangle = 1

class Config:
    """
    simulation parameter class
    """

    def __init__(self):
        # robot parameter
        self.max_speed = 0.1  # [m/s]
        self.min_speed = 0.02  # [m/s]

        self.max_accel = 0.02  # [m/ss]

        self.v_resolution = 0.01  # [m/s]
        self.yaw_rate_resolution = self.v_resolution / 5.0  # [rad/s]
        self.dt = 0.1  # [s] Time tick for motion prediction
        self.predict_time = 3  # [s]
        self.to_goal_cost_gain = 0.01
        self.speed_cost_gain = 1.0
        self.obstacle_cost_gain = 100
        self.robot_type = RobotType.rectangle

        # if robot_type == RobotType.circle
        # Also used to check if goal is reached in both types
        self.robot_radius = 0.65  # [m] for collision check

        # if robot_type == RobotType.rectangle
        self.robot_width = 1.3  # [m] for collision check
        self.robot_length = 2 # [m] for collision check


config = Config()

def motion(x, u, dt):
    """
    motion model
    u[0] -- 线速度
    u[1] -- 角速度

    x[2] -- 角度（rad)
    x[3] -- 速度（m/s）
    x[4] -- 角速度（rad/s）
    """

    x[2] += u[1] * dt
    x[0] += u[0] * math.cos(x[2]) * dt
    x[1] += u[0] * math.sin(x[2]) * dt
    x[3] = u[0]
    x[4] = u[1]

    return x


def calc_dynamic_window(x, config):
    """
    calculation dynamic window based on current state x
    """

    # Dynamic window from robot specification
    Vs = [config.min_speed, config.max_speed]

    # Dynamic window from motion model
    # 假设电机加速是瞬间的
    Vd = [x[3] - config.max_accel,
          x[3] + config.max_accel]

    #  [v_min, v_max, yaw_rate_min, yaw_rate_max]
    dw = [max(Vs[0], Vd[0]), min(Vs[1], Vd[1]),
          0, 0]

    dw[2] = -dw[1] / 5.0
    dw[3] = dw[1] / 5.0

    return dw


def predict_trajectory(x_init, v, y, config):
    """
    predict trajectory with an input
    """

    x = np.array(x_init)
    trajectory = np.array(x)
    time = 0
    while time <= config.predict_time:
        x = motion(x, [v, y], config.dt)
        trajectory = np.vstack((trajectory, x))
        time += config.dt

    return trajectory


def calc_control_and_trajectory(x, dw, config, goal, ob):
    """
    calculation final input with dynamic window
    """

    x_init = x[:]
    min_cost = float("inf")
    best_u = [0.0, 0.0]
    best_trajectory = np.array([x])

    #防止np.arange取不到数据，np.arange(0.1, 0.1, 0.1)返回为空数组
    err = 0.0001

    # evaluate all trajectory with sampled input in dynamic window
    for v in np.arange(dw[0], dw[1]+err, config.v_resolution):
        for y in np.arange(dw[2], dw[3]+err, config.yaw_rate_resolution):

            trajectory = predict_trajectory(x_init, v, y, config)

            # calc cost
            to_goal_cost = config.to_goal_cost_gain * calc_to_goal_cost(trajectory, goal)
            speed_cost = config.speed_cost_gain * (config.max_speed - trajectory[-1, 3])
            ob_cost = config.obstacle_cost_gain * calc_obstacle_cost(trajectory, ob, config)

            final_cost = to_goal_cost + speed_cost + ob_cost

            # search minimum trajectory
            if min_cost >= final_cost:
                min_cost = final_cost
                best_u = [v, y]
                best_trajectory = trajectory

    return best_u, best_trajectory


def calc_obstacle_cost(trajectory, ob, config):
    """
        calc obstacle cost inf: collision
    """
    ox = ob[:, 0]
    oy = ob[:, 1]
    dx = trajectory[:, 0] - ox[:, None]
    dy = trajectory[:, 1] - oy[:, None]
    r = np.hypot(dx, dy)

    if config.robot_type == RobotType.rectangle:
        yaw = trajectory[:, 2]
        rot = np.array([[np.cos(yaw), -np.sin(yaw)], [np.sin(yaw), np.cos(yaw)]])
        rot = np.transpose(rot, [2, 0, 1])
        local_ob = ob[:, None] - trajectory[:, 0:2]
        local_ob = local_ob.reshape(-1, local_ob.shape[-1])
        local_ob = np.array([local_ob @ x for x in rot])
        local_ob = local_ob.reshape(-1, local_ob.shape[-1])
        upper_check = local_ob[:, 0] <= config.robot_length / 2
        right_check = local_ob[:, 1] <= config.robot_width / 2
        bottom_check = local_ob[:, 0] >= -config.robot_length / 2
        left_check = local_ob[:, 1] >= -config.robot_width / 2
        if (np.logical_and(np.logical_and(upper_check, right_check),
                           np.logical_and(bottom_check, left_check))).any():
            return 10000
    elif config.robot_type == RobotType.circle:
        if np.array(r <= config.robot_radius).any():
            return 10000

    min_r = np.min(r)
    return 1.0 / min_r  # OK


def calc_to_goal_cost(trajectory, goal):
    """
        calc to goal cost with angle difference
    """

    dx = goal[0] - trajectory[-1, 0]
    dy = goal[1] - trajectory[-1, 1]
    error_angle = math.atan2(dy, dx)
    cost_angle = error_angle - trajectory[-1, 2]
    cost = abs(math.atan2(math.sin(cost_angle), math.cos(cost_angle)))

    return cost

class Simulation:

    def __init__(self):
        self.x = np.array(list(
            {
                "x":    0.0,
                "y":    0.0,
                "yaw":  0.0,
                "v":    0.0,
                "omega":0.0
            }.values()
        ))

    @staticmethod
    def plot_arrow(x, y, yaw, length=0.5, width=0.1):  # pragma: no cover
        plt.arrow(x, y, length * math.cos(yaw), length * math.sin(yaw),
                head_length=width, head_width=width)
        plt.plot(x, y)

    @staticmethod
    def plot_robot(x, y, yaw, config):  # pragma: no cover
        if config.robot_type == RobotType.rectangle:
            outline = np.array([[-config.robot_length / 2, config.robot_length / 2,
                                (config.robot_length / 2), -config.robot_length / 2,
                                -config.robot_length / 2],
                                [config.robot_width / 2, config.robot_width / 2,
                                - config.robot_width / 2, -config.robot_width / 2,
                                config.robot_width / 2]])
            Rot1 = np.array([[math.cos(yaw), math.sin(yaw)],
                            [-math.sin(yaw), math.cos(yaw)]])
            outline = (outline.T.dot(Rot1)).T
            outline[0, :] += x
            outline[1, :] += y
            plt.plot(np.array(outline[0, :]).flatten(),
                    np.array(outline[1, :]).flatten(), "-k")
        elif config.robot_type == RobotType.circle:
            circle = plt.Circle((x, y), config.robot_radius, color="b")
            plt.gcf().gca().add_artist(circle)
            out_x, out_y = (np.array([x, y]) +
                            np.array([np.cos(yaw), np.sin(yaw)]) * config.robot_radius)
            plt.plot([x, out_x], [y, out_y], "-k")

    def move(self, u, dt, goal, ob, predicted_trajectory):
        self.x = motion(self.x, u, dt)
        plt.cla()
        plt.axis("equal")
        plt.plot(predicted_trajectory[:, 0], predicted_trajectory[:, 1], "-g")
        plt.plot(self.x[0], self.x[1], "xr")
        plt.plot(goal[0], goal[1], "xb")
        plt.plot(ob[:, 0], ob[:, 1], "ok", markersize=1)
        Simulation.plot_robot(self.x[0], self.x[1], self.x[2], config)
        Simulation.plot_arrow(self.x[0], self.x[1], self.x[2])
        plt.grid(True)
        plt.pause(0.0001)

'''
模仿rospy的简易数据
'''
class _Twist:
    def __init__(self):
        self.vx = 0  #x方向速度
        self.vth = 0 #角速度

class _Point:
    def __init__(self):
        self.x = 0
        self.y = 0

class _Odom(_Point):
    def __init__(self):
        _Point.__init__(self)
        self.pose_z = 0  #角度


def uselock(f):
    scan_lock = Lock()
    def wrapper(self, *arg, **kargs):
        nonlocal scan_lock
        scan_lock.acquire()
        ret = f(self, *arg, **kargs)
        scan_lock.release()
        return ret
    return wrapper


class Core:
    def __init__(self):
        self.twist = _Twist()
        self.odom = _Odom()
        self.origin = _Point()
        self.goal = _Point()
        self.running = False
        self.get_goal = False #达成目标
        self.ob = None
        self.simulator = Simulation()

    def set_origin(self, *o):
        self.origin.x = o[0]
        self.origin.y = o[1]

    def set_goal(self, *g):
        self.goal.x = g[0]
        self.goal.y = g[1]

    @uselock
    def update_ob(self, scan):
        '''
        刷新障碍数据
        '''
        self.ob = get_obstacle_array(scan)

    @uselock
    def dwa_task(self, x, config, goal, ob):
        return dwa_control(x, config, goal, ob)

    def run_task(self, sim=False):
        self.running = True
        self.get_goal = False

        _x = np.array([
            self.odom.x, 
            self.odom.y, 
            self.odom.pose_z, 
            self.twist.vx, 
            self.twist.vth])

        _goal = np.array([self.goal.x, self.goal.y])

        tick = 0

        trajectory = np.array(_x)
        
        while self.running and self.ob is not None:

            tick+=1
            #time.sleep(config.dt)

            u, predict = self.dwa_task(_x, config, _goal, self.ob)
            trajectory = np.vstack((trajectory, _x))
            print("[%d] v: %.5f, angular: %.5f, x: %.5f, y: %.5f"%(tick, u[0], u[1], _x[0], _x[1]))

            if sim: 
                self.simulator.move(u, config.dt, _goal, self.ob, predict)
                _x = self.simulator.x

            #dist_to_goal = math.hypot(_x[0] - _goal[0], _x[1] - _goal[1])

            '''仅考虑x方向'''
            dist_to_goal = abs(_x[0] - _goal[0])

            if dist_to_goal <= 0.1:
                print("Goal!!")
                self.get_goal = True
                self.running = False

                plt.plot(trajectory[:, 0], trajectory[:, 1], "-r")
                plt.pause(0.0001)

                plt.show()
                break


